import os
import numpy as np
from PIL import Image

np.seterr(divide="ignore", invalid="ignore")


def loadDataSet(path, option):
    filePaths = []
    labels = []
    for subDir in os.listdir(path):
        filePath = os.path.join(path, subDir)
        if os.path.isdir(filePath):
            for fileName in os.listdir(filePath):
                if option == "train":
                    if eval(fileName[:-4]) <= 5:
                        fullPath = os.path.join(filePath, fileName)
                        filePaths.append(fullPath)
                        labels.append(filePath[6:])
                elif option == "test":
                    if eval(fileName[:-4]) >= 6:
                        fullPath = os.path.join(filePath, fileName)
                        filePaths.append(fullPath)
                        labels.append(filePath[6:])

    dataSet = []
    for file in filePaths:
        img = Image.open(file)
        dataSet.append(np.array(img))

    return np.array(dataSet), np.array(labels)


# 计算距离
def getDistance(vecA, vecB):
    return np.sqrt(np.sum(np.power(vecA - vecB, 2)))


# 初始化聚簇中心，取k个随机质心
def initCent(dataSet, k):
    centroids = []
    width, height = np.shape(dataSet[0])
    maxV = np.max(dataSet)
    minV = np.min(dataSet)
    rangeV = maxV - minV
    for i in range(k):
        cent = np.zeros([width, height]) + minV + rangeV * np.random.rand(width, height)
        centroids.append(cent)
    return np.array(centroids)


def train(dataSet, labels, k):
    centroids = initCent(dataSet, k)
    numOfSamples = np.shape(dataSet)[0]
    clusterAssment = np.zeros([numOfSamples, 2])

    clusterChanged = True
    while clusterChanged:
        clusterChanged = False
        for i in range(numOfSamples):
            minDist = 1e9
            minIndex = 0
            for j in range(k):
                distance = getDistance(centroids[j], dataSet[i])
                if distance < minDist:
                    minDist = distance
                    minIndex = j
            if clusterAssment[i, 0] != minIndex:
                clusterChanged = True
                clusterAssment[i, :] = minIndex, minDist ** 2
        for j in range(k):
            selected = np.nonzero(clusterAssment[:, 0] == j)[0]
            if len(selected):
                pointsInCluster = dataSet[np.nonzero(clusterAssment[:, 0] == j)[0]]
                centroids[j] = np.mean(pointsInCluster, axis=0)

    clusterJudge = [0 for i in range(40)]
    for i in range(40):
        clusterJudge[eval(labels[i][1:]) - 1] = np.argmax(np.bincount(np.array(clusterAssment[i:i + 5, 0]).astype(int)))
    clusterJudge = np.array(clusterJudge)

    return centroids, clusterJudge


def test(centroids, clusterStandard, dataSet, labels, k):
    numOfSamples = np.shape(dataSet)[0]

    clusterAssment = np.zeros([numOfSamples]).astype(int)

    for i in range(numOfSamples):
        minDist = 1e9
        minIndex = 0
        for j in range(k):
            distance = getDistance(centroids[j], dataSet[i])
            if distance < minDist:
                minDist = distance
                minIndex = j
        if clusterAssment[i] != minIndex:
            clusterAssment[i] = minIndex

    clusterJudge = []
    for label in labels:
        label = eval(label[1:])
        clusterJudge.append(clusterStandard[label - 1])
    clusterJudge = np.array(clusterJudge).astype(int)

    wrongPoints = np.nonzero(clusterAssment != clusterJudge)[0]

    wrongSamples = []
    for wp in wrongPoints:
        wrongSamples.append([labels[wp], wp % 5 + 6])

    accuracy = 1 - len(wrongSamples) / numOfSamples

    return accuracy, wrongSamples
